package com.lj.sse.service.impl;

import cn.dev33.satoken.stp.StpUtil;
import cn.hutool.core.thread.ThreadUtil;
import cn.hutool.core.util.IdUtil;
import com.lj.common_web.response.ResponseCode;
import com.lj.common_web.response.ResponseVO;
import com.lj.sse.properties.SseProperties;
import com.lj.sse.result.SseRecord;
import com.lj.sse.service.SseService;
import com.lj.sse.utils.SseContextUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;

/**
 * @author luojing
 * @since 2025/3/26 11:39
 */
@Service
@Slf4j
public class SseServiceImpl implements SseService {

    private final ScheduledThreadPoolExecutor executor;

    public SseServiceImpl(SseProperties sseProperties) {
        this.executor = ThreadUtil.createScheduledExecutor(sseProperties.getThreadNum());
    }

    @Override
    public SseEmitter createSseConnect(String clientId, long period, TimeUnit unit, Consumer<SseRecord> consumer) {
        long userId = StpUtil.getLoginIdAsLong();
        // 判断是否已经有连接
        SseRecord sseRecord = SseContextUtils.getSseRecord(userId, clientId);
        if (sseRecord != null) {
            return sseRecord.getEmitter();
        }
        // 创建一个新的sse连接
        SseEmitter sseEmitter = new SseEmitter(0L);
        SseRecord newSseRecord = new SseRecord(userId, IdUtil.fastSimpleUUID(), sseEmitter);
        if (consumer != null) {
            ScheduledFuture<?> scheduledFuture = executor.scheduleAtFixedRate(() -> consumer.accept(newSseRecord), 2, period, unit);
            newSseRecord.setScheduledFuture(scheduledFuture);
        }
        SseContextUtils.load(newSseRecord);
        // 向客户端发送客户端id
        SseContextUtils.send(newSseRecord, ResponseVO.of(ResponseCode.CREATED, newSseRecord.getClientId()));
        return sseEmitter;
    }

    @Override
    public void closeSseConnect(Long userId, String clientId) {
        SseContextUtils.unload(userId, clientId);
    }
}
